/**
 * WebGPU Batch Render Bindings - JSB 2.0 Style
 * 参考Cocos JSB 2.0创建的JS到C++绑定接口
 */

#pragma once

#include <memory>
#include <vector>
#include <unordered_map>
#include "dawn/webgpu.h"

// 前向声明
struct se_class_t;
struct se_object_t;
struct se_value_t;

namespace webgpu {

// 资源类型枚举
enum class ResourceType : uint8_t {
    Buffer = 0,
    Texture = 1,
    TextureView = 2,
    Sampler = 3,
    ShaderModule = 4,
    RenderPipeline = 5,
    ComputePipeline = 6,
    BindGroupLayout = 7,
    BindGroup = 8,
    PipelineLayout = 9,
    RenderBundle = 10,
    RenderBundleEncoder = 11,
    CommandEncoder = 12,
    CommandBuffer = 13
};

// 资源状态枚举
enum class ResourceState : uint8_t {
    Created = 0,
    InUse = 1,
    PendingDestroy = 2,
    Destroyed = 3
};

// 资源信息结构
struct ResourceInfo {
    uint32_t id;
    ResourceType type;
    ResourceState state;
    uint32_t refCount;
    size_t memorySize;
    std::chrono::steady_clock::time_point createTime;
    std::chrono::steady_clock::time_point lastAccessTime;
    std::vector<uint32_t> dependencies; // 依赖的其他资源ID
    std::vector<uint32_t> dependents;   // 被依赖的资源ID
    std::string label;

    ResourceInfo(uint32_t resId, ResourceType resType, const std::string& resLabel = "")
        : id(resId), type(resType), state(ResourceState::Created), refCount(1),
          memorySize(0), label(resLabel) {
        auto now = std::chrono::steady_clock::now();
        createTime = now;
        lastAccessTime = now;
    }
};

// 资源生命周期管理器
class ResourceLifecycleManager {
public:
    ResourceLifecycleManager();
    ~ResourceLifecycleManager();

    // 资源注册与注销
    void registerResource(uint32_t id, ResourceType type, const std::string& label = "");
    void unregisterResource(uint32_t id);

    // 引用计数管理
    void addReference(uint32_t id);
    void removeReference(uint32_t id);
    uint32_t getReferenceCount(uint32_t id) const;

    // 依赖关系管理
    void addDependency(uint32_t dependentId, uint32_t dependencyId);
    void removeDependency(uint32_t dependentId, uint32_t dependencyId);
    std::vector<uint32_t> getDependencies(uint32_t id) const;
    std::vector<uint32_t> getDependents(uint32_t id) const;

    // 内存管理
    void setMemorySize(uint32_t id, size_t size);
    size_t getTotalMemoryUsage() const;
    size_t getMemoryUsageByType(ResourceType type) const;

    // 资源状态管理
    void setState(uint32_t id, ResourceState state);
    ResourceState getState(uint32_t id) const;
    void updateLastAccessTime(uint32_t id);

    // 清理操作
    void cleanup();
    void cleanupByType(ResourceType type);
    void forceDestroy(uint32_t id);
    std::vector<uint32_t> getResourcesForCleanup() const;

    // 验证与调试
    bool isValid(uint32_t id) const;
    ResourceInfo* getResourceInfo(uint32_t id);
    std::vector<ResourceInfo> getAllResources() const;
    void dumpResourceInfo() const;

    // 配置
    void setMaxMemoryUsage(size_t maxBytes);
    void setCleanupThreshold(size_t count);
    void setAutoCleanupEnabled(bool enabled);

private:
    std::unordered_map<uint32_t, ResourceInfo> resources_;
    std::mutex resourceMutex_;

    // 配置参数
    size_t maxMemoryUsage_;
    size_t cleanupThreshold_;
    bool autoCleanupEnabled_;

    // 统计信息
    mutable std::mutex statsMutex_;
    std::unordered_map<ResourceType, size_t> memoryByType_;

    // 内部方法
    void destroyResourceInternal(uint32_t id);
    bool canDestroy(uint32_t id) const;
    void updateDependencies(uint32_t id);
    void checkMemoryPressure();
};

// 资源池管理器
template<typename T>
class ResourcePool {
public:
    ResourcePool(size_t maxSize = 100) : maxSize_(maxSize) {}

    std::shared_ptr<T> acquire() {
        std::lock_guard<std::mutex> lock(mutex_);
        if (!pool_.empty()) {
            auto resource = pool_.back();
            pool_.pop_back();
            return resource;
        }
        return nullptr;
    }

    void release(std::shared_ptr<T> resource) {
        std::lock_guard<std::mutex> lock(mutex_);
        if (pool_.size() < maxSize_) {
            pool_.push_back(resource);
        }
    }

    void clear() {
        std::lock_guard<std::mutex> lock(mutex_);
        pool_.clear();
    }

private:
    std::vector<std::shared_ptr<T>> pool_;
    std::mutex mutex_;
    size_t maxSize_;
};

// 命令类型枚举，与JS端保持一致
enum class CommandType : uint8_t {
    // Device & Adapter
    CREATE_DEVICE = 0x01,
    DESTROY_DEVICE = 0x02,

    // Buffer Operations
    CREATE_BUFFER = 0x10,
    DESTROY_BUFFER = 0x11,
    WRITE_BUFFER = 0x12,
    COPY_BUFFER = 0x13,

    // Texture Operations
    CREATE_TEXTURE = 0x20,
    DESTROY_TEXTURE = 0x21,
    WRITE_TEXTURE = 0x22,
    COPY_TEXTURE = 0x23,

    // Render Pipeline
    CREATE_RENDER_PIPELINE = 0x30,
    DESTROY_RENDER_PIPELINE = 0x31,
    CREATE_SHADER_MODULE = 0x32,
    DESTROY_SHADER_MODULE = 0x33,

    // Bind Group & Layout
    CREATE_BIND_GROUP_LAYOUT = 0x34,
    DESTROY_BIND_GROUP_LAYOUT = 0x35,
    CREATE_BIND_GROUP = 0x36,
    DESTROY_BIND_GROUP = 0x37,
    CREATE_PIPELINE_LAYOUT = 0x38,
    DESTROY_PIPELINE_LAYOUT = 0x39,

    // Texture View & Sampler
    CREATE_TEXTURE_VIEW = 0x3A,
    DESTROY_TEXTURE_VIEW = 0x3B,
    CREATE_SAMPLER = 0x3C,
    DESTROY_SAMPLER = 0x3D,

    // Compute Pipeline
    CREATE_COMPUTE_PIPELINE = 0x40,
    DESTROY_COMPUTE_PIPELINE = 0x41,

    // Command Encoder
    CREATE_COMMAND_ENCODER = 0x50,
    FINISH_COMMAND_ENCODER = 0x51,

    // Render Pass
    BEGIN_RENDER_PASS = 0x60,
    END_RENDER_PASS = 0x61,
    SET_PIPELINE = 0x62,
    SET_VERTEX_BUFFER = 0x63,
    SET_INDEX_BUFFER = 0x64,
    SET_BIND_GROUP = 0x65,
    DRAW = 0x66,
    DRAW_INDEXED = 0x67,

    // Compute Pass
    BEGIN_COMPUTE_PASS = 0x70,
    END_COMPUTE_PASS = 0x71,
    DISPATCH = 0x72,

    // Render Bundle
    CREATE_RENDER_BUNDLE_ENCODER = 0x73,
    FINISH_RENDER_BUNDLE_ENCODER = 0x74,
    EXECUTE_BUNDLES = 0x75,

    // Queue Operations
    SUBMIT = 0x80,

    // Batch Control
    BATCH_START = 0xF0,
    BATCH_END = 0xF1,
    FLUSH = 0xFF
};

// 命令头部结构
#pragma pack(push, 1)
struct CommandHeader {
    uint8_t type;
    uint8_t flags;
    uint16_t id;
    uint32_t size;
};
#pragma pack(pop)

// 批量渲染器类
class BatchRenderer {
public:
    BatchRenderer();
    ~BatchRenderer();
    
    // 初始化Dawn
    bool initialize();
    
    // 处理批量命令
    void processBatch(const uint8_t* data, size_t size);
    
    // 获取设备
    WGPUDevice getDevice() const { return device_; }

private:
    WGPUInstance instance_;
    WGPUAdapter adapter_;
    WGPUDevice device_;
    WGPUQueue queue_;

    // 资源管理
    std::unordered_map<uint32_t, WGPUBuffer> buffers_;
    std::unordered_map<uint32_t, WGPUTexture> textures_;
    std::unordered_map<uint32_t, WGPUTextureView> textureViews_;
    std::unordered_map<uint32_t, WGPUSampler> samplers_;
    std::unordered_map<uint32_t, WGPUShaderModule> shaderModules_;
    std::unordered_map<uint32_t, WGPURenderPipeline> renderPipelines_;
    std::unordered_map<uint32_t, WGPUComputePipeline> computePipelines_;
    std::unordered_map<uint32_t, WGPUBindGroupLayout> bindGroupLayouts_;
    std::unordered_map<uint32_t, WGPUBindGroup> bindGroups_;
    std::unordered_map<uint32_t, WGPUPipelineLayout> pipelineLayouts_;
    std::unordered_map<uint32_t, WGPURenderBundle> renderBundles_;
    std::unordered_map<uint32_t, WGPURenderBundleEncoder> renderBundleEncoders_;
    std::unordered_map<uint32_t, WGPUCommandEncoder> commandEncoders_;

    // 当前状态
    WGPUCommandEncoder currentEncoder_;
    WGPURenderPassEncoder currentRenderPass_;
    WGPUComputePassEncoder currentComputePass_;
    std::vector<WGPUCommandBuffer> pendingCommandBuffers_;

    // 资源生命周期管理
    std::unique_ptr<ResourceLifecycleManager> lifecycleManager_;

    // 资源池
    ResourcePool<WGPUBuffer> bufferPool_;
    ResourcePool<WGPUTexture> texturePool_;
    
    // 命令处理
    void processCommand(const uint8_t*& data, size_t& remaining);
    
    // 具体命令处理函数
    void handleCreateBuffer(const uint8_t*& data, size_t& remaining);
    void handleDestroyBuffer(const uint8_t*& data, size_t& remaining);
    void handleWriteBuffer(const uint8_t*& data, size_t& remaining);
    void handleCreateTexture(const uint8_t*& data, size_t& remaining);
    void handleCreateTextureView(const uint8_t*& data, size_t& remaining);
    void handleCreateSampler(const uint8_t*& data, size_t& remaining);
    void handleCreateShaderModule(const uint8_t*& data, size_t& remaining);
    void handleCreateRenderPipeline(const uint8_t*& data, size_t& remaining);
    void handleCreateComputePipeline(const uint8_t*& data, size_t& remaining);
    void handleCreateBindGroupLayout(const uint8_t*& data, size_t& remaining);
    void handleCreateBindGroup(const uint8_t*& data, size_t& remaining);
    void handleCreatePipelineLayout(const uint8_t*& data, size_t& remaining);
    void handleCreateRenderBundleEncoder(const uint8_t*& data, size_t& remaining);
    void handleFinishRenderBundleEncoder(const uint8_t*& data, size_t& remaining);
    void handleWriteTexture(const uint8_t*& data, size_t& remaining);
    void handleCopyBufferToBuffer(const uint8_t*& data, size_t& remaining);
    void handleCopyTextureToTexture(const uint8_t*& data, size_t& remaining);
    void handleBeginRenderPass(const uint8_t*& data, size_t& remaining);
    void handleEndRenderPass(const uint8_t*& data, size_t& remaining);
    void handleBeginComputePass(const uint8_t*& data, size_t& remaining);
    void handleEndComputePass(const uint8_t*& data, size_t& remaining);
    void handleSetPipeline(const uint8_t*& data, size_t& remaining);
    void handleSetVertexBuffer(const uint8_t*& data, size_t& remaining);
    void handleSetIndexBuffer(const uint8_t*& data, size_t& remaining);
    void handleSetBindGroup(const uint8_t*& data, size_t& remaining);
    void handleDraw(const uint8_t*& data, size_t& remaining);
    void handleDrawIndexed(const uint8_t*& data, size_t& remaining);
    void handleExecuteBundles(const uint8_t*& data, size_t& remaining);

    // Render Pass 扩展方法
    void handleSetViewport(const uint8_t*& data, size_t& remaining);
    void handleSetScissorRect(const uint8_t*& data, size_t& remaining);
    void handleSetBlendConstant(const uint8_t*& data, size_t& remaining);
    void handleSetStencilReference(const uint8_t*& data, size_t& remaining);

    // Render Bundle 命令处理函数
    void handleBundleSetPipeline(const uint8_t*& data, size_t& remaining);
    void handleBundleSetVertexBuffer(const uint8_t*& data, size_t& remaining);
    void handleBundleSetIndexBuffer(const uint8_t*& data, size_t& remaining);
    void handleBundleSetBindGroup(const uint8_t*& data, size_t& remaining);
    void handleBundleDraw(const uint8_t*& data, size_t& remaining);
    void handleBundleDrawIndexed(const uint8_t*& data, size_t& remaining);

    void handleDispatch(const uint8_t*& data, size_t& remaining);
    void handleCreateCommandEncoder(const uint8_t*& data, size_t& remaining);
    void handleFinishCommandEncoder(const uint8_t*& data, size_t& remaining);
    void handleSubmit(const uint8_t*& data, size_t& remaining);
    
    // 辅助函数
    uint32_t readUInt32(const uint8_t*& data, size_t& remaining);
    int32_t readInt32(const uint8_t*& data, size_t& remaining);
    uint16_t readUInt16(const uint8_t*& data, size_t& remaining);
    uint8_t readUInt8(const uint8_t*& data, size_t& remaining);
    float readFloat32(const uint8_t*& data, size_t& remaining);
    std::string readString(const uint8_t*& data, size_t& remaining);
    std::vector<uint8_t> readBuffer(const uint8_t*& data, size_t& remaining);

    // Dawn 枚举反序列化函数
    WGPULoadOp deserializeLoadOp(uint32_t value);
    WGPUStoreOp deserializeStoreOp(uint32_t value);
    WGPUCompareFunction deserializeCompareFunction(uint32_t value);
    WGPUStencilOperation deserializeStencilOperation(uint32_t value);
    WGPUPrimitiveTopology deserializePrimitiveTopology(uint32_t value);
    WGPUFrontFace deserializeFrontFace(uint32_t value);
    WGPUCullMode deserializeCullMode(uint32_t value);
    WGPUBlendOperation deserializeBlendOperation(uint32_t value);
    WGPUBlendFactor deserializeBlendFactor(uint32_t value);
    WGPUAddressMode deserializeAddressMode(uint32_t value);
    WGPUFilterMode deserializeFilterMode(uint32_t value);
    WGPUMipmapFilterMode deserializeMipmapFilterMode(uint32_t value);
    WGPUTextureFormat deserializeTextureFormat(uint32_t value);
    WGPUBufferBindingType deserializeBufferBindingType(uint32_t value);
    WGPUSamplerBindingType deserializeSamplerBindingType(uint32_t value);
    WGPUTextureSampleType deserializeTextureSampleType(uint32_t value);
    WGPUTextureViewDimension deserializeTextureViewDimension(uint32_t value);
    WGPUStorageTextureAccess deserializeStorageTextureAccess(uint32_t value);
    WGPUTextureAspect deserializeTextureAspect(uint32_t value);
    WGPUTextureDimension deserializeTextureDimension(uint32_t value);
    WGPUIndexFormat deserializeIndexFormat(uint32_t value);
    
    static void onDeviceError(WGPUErrorType type, const char* message, void* userdata);

    // 资源生命周期管理方法
    void registerResource(uint32_t id, ResourceType type, const std::string& label = "");
    void addResourceDependency(uint32_t dependentId, uint32_t dependencyId);
    void setResourceMemorySize(uint32_t id, ResourceType type, size_t size);
    void destroyResource(uint32_t id, ResourceType type);
    void performResourceCleanup();

    // 资源池管理
    template<typename T> std::shared_ptr<T> acquireFromPool();
    template<typename T> void releaseToPool(std::shared_ptr<T> resource);
};

// WebGPU资源finalize回调数据
struct WebGPUResourceData {
    uint32_t resourceId;
    ResourceType resourceType;
    BatchRenderer* renderer;

    WebGPUResourceData(uint32_t id, ResourceType type, BatchRenderer* r)
        : resourceId(id), resourceType(type), renderer(r) {}
};

// JSB 2.0 绑定类
class JSB_WebGPUBatch {
public:
    static bool init(se_object_t* obj);
    static void destroy();

    static bool webgpuBatchRender(se_state_t& s);
    static bool webgpuGetDevice(se_state_t& s);
    static bool webgpuSetBatchMode(se_state_t& s);

    // 创建带finalize回调的WebGPU对象
    static se_object_t* createResourceObject(uint32_t id, ResourceType type, const char* className);

private:
    static std::unique_ptr<BatchRenderer> s_batchRenderer;
    static se_class_t* s_class;
    static se_object_t* s_proto;

    // 各种资源类型的SE类
    static se_class_t* s_bufferClass;
    static se_class_t* s_textureClass;
    static se_class_t* s_pipelineLayoutClass;
    static se_class_t* s_bindGroupClass;
    static se_class_t* s_samplerClass;
    static se_class_t* s_renderPipelineClass;
    static se_class_t* s_computePipelineClass;

    SE_DECLARE_FINALIZE_FUNC(JSB_WebGPUBatch_finalize);
    SE_DECLARE_FINALIZE_FUNC(JSB_WebGPUResource_finalize);

    // 初始化资源类
    static void initResourceClasses(se_object_t* global);
};

} // namespace webgpu

// C风格导出函数，供JSB使用
extern "C" {
    bool jsb_register_webgpu_batch(se_object_t* obj);
}